#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/5/15 11:12
# @Author  : CHEN Wang
# @Site    :
# @File    : my_excel.py
# @Software: PyCharm

"""
脚本说明: 对Mongo数据库的连接，及各种操作
"""


import time
import pandas as pd
from pymongo import MongoClient
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.conf import LOCAL_CONF
from quant_researcher.quant.project_tool import assert_tool
from quant_researcher.quant.project_tool import hammer
import pymongo
from motor.motor_asyncio import AsyncIOMotorClient
# from motor import MotorClient
import asyncio

MichaelJ = None  # 用来连接行情表的专用


def get_updated_on():
    return time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))


def write_df(df, db=None, table=None, p_keys=None, **kwargs):
    """
    保存 DF 到 MongoDB

    :param DF df: 将要存储的数据
    :param str db: 数据库，默认：None（将从配置读取）
    :param str table: 表，默认：None（当然不可以这样，只是因为历史原因，要不然把它放前面了）
    :param str或list p_keys: 主键，默认：None
    :param kwargs:
        - set_updated_on，是不是设置更新时间，默认：False
    :return:
    """
    set_updated_on = kwargs.get('set_updated_on', False)
    if db is None:
        db = LOCAL_CONF['mongo']['db_in_use']

    if isinstance(p_keys, str):
        p_keys = [p_keys]
    LOG.info(f'去往的表（主键：{p_keys}）和库：%s @ %s，是否增加更新时间：{set_updated_on}→{bool(set_updated_on)}。', table, db)

    if set_updated_on:
        df['updated_on'] = get_updated_on()

    client = get_client()
    my_cols = df.columns.values
    tmp_range = range(len(my_cols))
    try:
        conn = client[db][table]
        for ii in df.values:
            tmp_doc = {}
            filter_dict = {}
            for jj in tmp_range:
                this_col = my_cols[jj]
                if p_keys is not None:
                    if this_col in p_keys:
                        filter_dict.update({this_col: ii[jj]})
                    else:
                        if tmp_doc.__contains__('$set'):
                            tmp_doc['$set'].update({this_col: ii[jj]})
                        else:
                            tmp_doc.update({'$set': {this_col: ii[jj]}})
                else:
                    tmp_doc.update({this_col: ii[jj]})
            if p_keys is None:
                conn.insert(tmp_doc)
            else:
                # LOG.info('filter: %s', filter_dict)
                # LOG.info('set: %s', tmp_doc)
                # assert False
                result = conn.update_one(filter_dict, tmp_doc, upsert=True)
                if result.matched_count > 1:
                    LOG.error('找到记录%s个，更新记录%s个', result.matched_count, result.modified_count)
                    LOG.error('找到的匹配项不唯一，请检查！！！')
                    raise RuntimeError('找到的匹配项不唯一')
    except Exception as err:
        LOG.error(err)
        raise err
    finally:
        client.close()


def update_one(res_db, res_table, document, primary_key_lst, set_list=None, **kwargs):
    """
    按照主键更新 MongoDB 中的数据
        单条或多条记录的单个或多个字段，且可以自动增加更新时间列
    :param str res_db: 目标库
    :param str res_table: 目标表
    :param dict或list document: a dict or a list of dicts，数据
    :param list primary_key_lst: 主键列表，按照此列表进行更新操作
    :param list set_list: 需要更新的字段，默认：None
    :param kwargs:
        - set_updated_on，bool，是不是设置一个更新时间在数据上，默认：False
        - upsert，bool，是不是在对应主键记录不存在的时候插入数据，默认：True
        - mongo_clt，mongo_clt，mongo_clt，默认：None
    :return: None
    """
    mongo_clt = kwargs.pop('mongo_clt', None)
    upsert = kwargs.pop('upsert', True)
    set_updated_on = kwargs.get('set_updated_on', False)  # 这里被调用的时，都尽量传递 bool，别用些还需要猜的东西
    LOG.info(f'set_updated_on={set_updated_on}, upsert={upsert}')

    document = [document] if not isinstance(document, list) else document

    if set_list is None:  # 寻找需要更新哪些字段值
        set_list = []
        for ii in document[0]:  # 如果没传，则更新除主键以外的所有字段
            if ii not in primary_key_lst:
                set_list.append(ii)

    close_this_clt = False
    if mongo_clt is None:
        mongo_clt = get_client()
        close_this_clt = True

    num_matched = num_modified = num_inserted = 0

    def core(x):
        mark_dict = {}
        for key in primary_key_lst:  # 去数据中把主键的值取出来，等等根据这些主键值进行更新
            mark_dict.update({key: x[key]})

        set_dict = {}
        for key in set_list:
            set_dict.update({key: x[key]})
        if set_updated_on:
            set_dict.update({'updated_on': get_updated_on()})
        set_dict = {'$set': set_dict}

        # hammer.slim(x, a_f=1, name='x')
        # hammer.slim(mark_dict, a_f=1, name='condition')
        # hammer.slim(set_dict, a_f=1, name='value')
        res = mongo_clt[res_db][res_table].update_one(mark_dict, set_dict, upsert=upsert)
        if res.matched_count > 1:
            LOG.error('数据：%s', x)
            raise RuntimeError(f'根据主键（{primary_key_lst}）找到多条（{res.matched_count}）记录')
        return res.matched_count, res.modified_count, (1 if res.upserted_id else 0)

    try:
        for y in document:
            ans = core(y)
            num_matched += ans[0]
            num_modified += ans[1]
            num_inserted += ans[2]
        LOG.info(f'找到 {num_matched}，更新 {num_modified}，新增 {num_inserted}')
    except Exception as err:
        LOG.error(f'尝试更新时，遇到错误：{str(err)}，下面尝试打印引起错误的数据')
        hammer.slim(primary_key_lst, name='主键列表')
        hammer.slim(set_list, name='更新列')
        hammer.slim(document, name='引起错误的数据')
        raise err
    finally:
        if close_this_clt:
            LOG.warning('Going to ♻ this MongoDB client created by myself!!!')
            mongo_clt.close()


def get_mongo_info(prod_or_test):
    if prod_or_test == 'prod':
        host = LOCAL_CONF['mongo']['host_prod']
        port = LOCAL_CONF.getint('mongo', 'port_prod')
    elif prod_or_test == 'test':
        host = LOCAL_CONF['mongo']['host_test']
        port = LOCAL_CONF.getint('mongo', 'port_test')
    else:
        raise NotImplementedError
    return host, port


def get_client(host='', port=27017, prod_or_test='prod', **kwargs):
    username = kwargs.pop('username', None)
    password = kwargs.pop('password', None)
    database = kwargs.pop('database', None)
    if username and password:
        if host == '' or host is None:
            LOG.warning('未传入MongoDB的host信息或传入的host是None，调用函数获得')
            host, port = get_mongo_info(prod_or_test)
        LOG.warning(f'即将创建到MongoDB的client，host={host}，port={port}')
        mongo_uri = 'mongodb://%s:%s@%s:%s' % (username, password, host, port)
        client = MongoClient(mongo_uri)
        # client = MongoClient(host, port)
    else:
        if host == '' or host is None:
            LOG.warning('未传入MongoDB的host信息或传入的host是None，调用函数获得')
            host, port = get_mongo_info(prod_or_test)
        LOG.warning(f'即将创建到MongoDB的client，host={host}，port={port}')
        client = MongoClient(host, port)

    if database:
        return client[database]
    else:
        return client


def connect_mongo(host, port, username, password, db):
    if username and password:
        mongo_uri = 'mongodb://%s:%s@%s:%s/%s' % (username, password, host, port, db)
        conn = MongoClient(mongo_uri)
    else:
        conn = MongoClient(host, port)
    return conn[db]


def get_hq_conn(db, table):
    global MichaelJ
    if db is None or table is None:
        err_msg = '传入的db或table是None'
        LOG.error(err_msg)
        raise ValueError(err_msg)
    if MichaelJ is None:
        LOG.info('第一次建立连接')
        client = get_client()
        MichaelJ = client[db][table]
    else:
        pass
        # LOG.info('已经存在此库-表的连接，直接返回')
    return MichaelJ


def find(filter_dict, table, **kwargs):
    """
    查询 MongoDB 的函数

    :param dict filter_dict: 查询条件字典
    :param str table: 查询哪一张表
    :param kwargs:
        - rename_these，str 或者 list of str，将这（些）列进行重命名，默认：[]
        - host，str，主机IP，默认：None
        - limit，int，数据条数，默认：None
        - res_type，str，数据对象类型，默认：df
        - sort_key，排序列，默认：None
            可取值：
                end_date，str
                [（’end_date‘， 1）]，list of tuples，其中 1：小到大；-1：大到小
        - sort_direction，排序方向，默认：None
            注意：
                只有在 sort_key 是字符串的时候，才能指定这个值
                如果 sort_key 是 list 的话，需要在 list 中通过 tuple 指定排序方向
        - db，str，使用到的数据库，默认：mongo 配置中的 db_in_use
        - projection_dict，dict，需要那些列或者不需要那些列的指定字典，默认：{'_id': False}
    :return: pd.DataFrame
    """
    rename_these = kwargs.pop('rename_these', [])
    host = kwargs.pop('host', None)
    limit = kwargs.pop('limit', None)
    res_type = kwargs.pop('res_type', 'df')
    sort_key = kwargs.pop('sort_key', None)
    sort_direction = kwargs.pop('sort_direction', None)
    db = kwargs.pop('db', LOCAL_CONF['mongo']['db_in_use'])
    projection_dict = kwargs.pop('projection_dict', {})
    projection_dict.update({'_id': False})
    user = kwargs.pop('user', None)
    password = kwargs.pop('password', None)
    prod_or_test = kwargs.pop('prod_or_test', 'prod')

    rename_these = [rename_these] if isinstance(rename_these, str) else rename_these
    assert_tool.value_in_list(res_type, ['df', 'list'])
    client = get_client(host, prod_or_test)

    try:
        conn = client[db]
        if user is not None and password is not None:
            conn.authenticate(user, password)
        conn = conn[table]
        cursor = conn.find(filter_dict, projection_dict, **kwargs)

        if sort_key is not None:
            if sort_direction is None:
                sort_direction = 1
            if isinstance(sort_key, str):
                sort_key = [(sort_key, sort_direction)]
            LOG.info('key or list for sorting: %s', sort_key)
            LOG.info('sort-direction for sorting: %s', sort_direction)
            cursor = cursor.sort(key_or_list=sort_key)
        else:
            if sort_direction is not None:
                err_msg = f'输入了排序方向（sort_direction={sort_direction}），' \
                          f'却没有指定字段，函数参数：{locals()}'
                LOG.error(err_msg)
                raise RuntimeError(err_msg)

        if limit is not None:
            cursor = cursor.limit(limit=limit)

        data = list(cursor)
        if res_type == 'df':
            LOG.warning(f'从MongoDB取到 {len(data)} 条数据，'
                        f'条件={filter_dict}，列指定={projection_dict}')
            data = pd.DataFrame(data)
            columns = {}
            for a_col in rename_these:
                before, after = a_col.split(' as ')
                columns.update({before: after})
            data.rename(columns=columns, inplace=True) if columns else None
            hammer.slim(
                data, a_f=1,
                name=f'从MongoDB取到的数据，条件={filter_dict}，列指定={projection_dict}'
            )
        if res_type == 'list' and limit == 1:
            # 如果返回结果为list，且limit为1，这又是MongoDB，那么99%的情况都会接下来再接一个[0]，那么这里直接先这样做了
            data = data[0]
        return data
    finally:
        client.close()
        LOG.info('This mongo client has been closed')


def c_find(columns, no_id=True):
    projection = {column: 1 for column in columns}
    if no_id:
        projection['_id'] = 0
    return projection


def c_range(column, start, end, right=True, left=True):
    l_sign = '$gte' if left else '$gt'
    r_sign = '$lte' if right else '$lt'
    return {column: {l_sign: start, r_sign: end}}


def db_filter(column, start, end, right=True, left=True, **kwargs):
    col_range = c_range(column, start, end, right=right, left=left)
    for key, value in kwargs.items():
        if isinstance(value, list):
            col_range[key] = {'$in': value}
        else:
            col_range[key] = value
    return col_range


def QA_util_sql_mongo_setting(uri='mongodb://localhost:27017/quantaxis'):
    """
    explanation:
        根据给定的uri返回一个MongoClient实例，采用@几何建议以使用加密

    params:
        * uri ->:
            meaning: mongodb连接uri
            type: str
            optional: [null]

    return:
        MongoClient

    demonstrate:
        Not described

    output:
        Not described
    """

    # 采用@几何的建议,使用uri代替ip,port的连接方式
    # 这样可以对mongodb进行加密:
    # uri=mongodb://user:passwor@ip:port
    client = pymongo.MongoClient(uri)
    return client


# def QA_util_sql_mongo_setting(host, port, user, password, db):
#     """
#     explanation:
#         根据给定的uri返回一个MongoClient实例，采用@几何建议以使用加密
#
#     params:
#         * uri ->:
#             meaning: mongodb连接uri
#             type: str
#             optional: [null]
#
#     return:
#         MongoClient
#
#     demonstrate:
#         Not described
#
#     output:
#         Not described
#     """
#
#     # 采用@几何的建议,使用uri代替ip,port的连接方式
#     # 这样可以对mongodb进行加密:
#     # uri=mongodb://user:passwor@ip:port
#     client = pymongo.MongoClient(host, port)
#     client = client[db]
#     client.authenticate(user, password)
#     return client
#
# # async


def QA_util_sql_async_mongo_setting(uri):
    """
    explanation:
        根据给定的uri返回一个异步AsyncIOMotorClient实例

    params:
        * uri ->:
            meaning: mongodb连接uri
            type: str
            optional: [null]

    return:
        AsyncIOMotorClient

    demonstrate:
        Not described

    output:
        Not described
    """
    # loop = asyncio.new_event_loop()
    # asyncio.set_event_loop(loop)
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

    # async def client():
    return AsyncIOMotorClient(uri, io_loop=loop)
    # yield  client()


ASCENDING = pymongo.ASCENDING
DESCENDING = pymongo.DESCENDING
QA_util_sql_mongo_sort_ASCENDING = pymongo.ASCENDING
QA_util_sql_mongo_sort_DESCENDING = pymongo.DESCENDING


if __name__ == '__main__':
    get_client(host='', port='', username=None, password=None, database='info_dbv1')

    # test async_mongo
    client = QA_util_sql_async_mongo_setting().quantaxis.stock_day
    print(client)